import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from src.util.save_mp4 import save_mp4
from src.util.trajectory_manager import TrajectoryManager


class SystemSimulation:
    def __init__(self, system, bodies, title="Celestial System Simulation"):
        self.system = system
        self.bodies = bodies if isinstance(bodies, list) else [bodies]
        self.title = title

        # 使用统一的轨迹管理器
        self.trajectory_manager = TrajectoryManager(max_trace_points=100, max_trail_points=30)

        # 设置图形
        self.fig, self.ax = plt.subplots(figsize=(8, 8))
        self.setup_plot()

        # 创建绘制对象
        self.body_plots = {}
        self.trace_plots = {}
        self.trail_plots = {}
        self.create_plots()

    def setup_plot(self):
        """设置图表样式"""
        # 设置背景为黑色
        self.fig.set_facecolor('black')
        self.ax.set_facecolor('black')
        # 设置坐标轴范围
        self.ax.set_xlim(-3, 3)
        self.ax.set_ylim(-3, 3)
        # 设置坐标轴比例
        self.ax.set_aspect('equal')
        # 设置标题颜色为白色
        self.ax.set_title(self.title, color='white')

        # 设置坐标轴标签颜色为白色
        self.ax.tick_params(colors='white')
        for spine in self.ax.spines.values():
            spine.set_color('white')
        self.ax.xaxis.label.set_color('white')
        self.ax.yaxis.label.set_color('white')

    def create_plots(self):
        """创建天体和轨迹的绘制对象"""
        for body in self.bodies:
            # 为每个天体初始化轨迹跟踪
            self.trajectory_manager.add_body(body.name)

            # 天体绘制对象
            self.body_plots[body.name], = self.ax.plot(
                [], [], 'o', color=body.color, markersize=body.radius, label=body.name
            )
            # 轨道线绘制对象
            self.trace_plots[body.name], = self.ax.plot(
                [], [], '-', color=body.color, alpha=0.7, linewidth=1
            )
            # 拖尾线绘制对象
            self.trail_plots[body.name], = self.ax.plot(
                [], [], '-', color=body.color, alpha=0.3, linewidth=0.5
            )

    def update_animation(self, frame):
        """动画更新函数"""
        # 系统演化
        self.system.evolve(0.1)

        # 更新每个天体的显示
        updated_plots = []
        for body in self.bodies:
            # 更新轨迹数据
            self.trajectory_manager.update_position(
                body.name, body.position[0], body.position[1]
            )

            # 更新天体位置
            plot = self.body_plots[body.name]
            plot.set_data([body.position[0]], [body.position[1]])
            updated_plots.append(plot)

            # 更新完整轨道
            trace_x, trace_y = self.trajectory_manager.get_trace_data(body.name)
            trace_plot = self.trace_plots[body.name]
            trace_plot.set_data(trace_x, trace_y)
            updated_plots.append(trace_plot)

            # 更新拖尾效果
            trail_x, trail_y = self.trajectory_manager.get_trail_data(body.name)
            if len(trail_x) > 1:
                trail_plot = self.trail_plots[body.name]
                trail_plot.set_data(trail_x, trail_y)
                updated_plots.append(trail_plot)

        self.ax.legend()
        return updated_plots

    def run_simulation(self, frames=200, interval=50, output_file=None):
        """运行模拟动画"""
        # 创建动画对象
        anim = FuncAnimation(
            self.fig, self.update_animation, frames=frames,
            interval=interval, blit=True
        )

        # 保存为MP4文件（如果指定了输出文件）
        if output_file:
            # 使用统一的save_mp4函数保存
            save_mp4(anim,
                     output_file=output_file,
                     fps=int(1000 / interval),  # 根据interval计算合适的fps
                     )
        # 显示动画
        plt.show()
